package cgl.narada.service.security.kmc;

import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.cert.Certificate;
import java.util.Calendar;
import java.util.Hashtable;
import java.util.Properties;
import java.util.Random;

import org.apache.log4j.Level;
import org.apache.log4j.Logger;

import cgl.narada.discovery.broker.BrokerDiscoveryHelper;
import cgl.narada.discovery.broker.messages.BrokerDiscoveryResponse;
import cgl.narada.discovery.topics.TopicDiscoveryClient;
import cgl.narada.discovery.topics.Topics;
import cgl.narada.discovery.topics.messages.EncryptedMessagePayload;
import cgl.narada.discovery.topics.messages.SignedTopicAdvertisement;
import cgl.narada.event.NBEvent;
import cgl.narada.event.TemplateProfileAndSynopsisTypes;
import cgl.narada.matching.Profile;
import cgl.narada.service.ServiceException;
import cgl.narada.service.client.ClientService;
import cgl.narada.service.client.EventConsumer;
import cgl.narada.service.client.EventProducer;
import cgl.narada.service.client.NBEventListener;
import cgl.narada.service.client.SessionService;
import cgl.narada.service.security.kmc.messages.KMCMessage;
import cgl.narada.service.security.kmc.messages.SecureTopicKeyCreateUpdateRequest;
import cgl.narada.service.security.kmc.messages.SecureTopicKeyRequest;
import cgl.narada.service.security.kmc.messages.SecureTopicKeyResponse;
import cgl.narada.service.security.kmc.messages.SignedRequest;

/**
 * <br>
 * Created on Sep 9, 2005
 * 
 * @author Harshawardhan Gadgil (hgadgil@grids.ucs.indiana.edu)
 */
public class KMCClient implements NBEventListener {

    static Logger log = Logger.getLogger("KMCClient");
    /*static {
        Logger.getRootLogger().setLevel(Level.WARN);
    }*/

    public static final byte MSG_SECURITY_TOKEN = 1;
    public static final byte MSG_ERROR = 2;

    private String kmcClientTopic;

    private Certificate myCertificate;
    private PrivateKey myPrivateKey;
    // private PublicKey rootCAPublicKey;

    // Util
    private TopicDiscoveryClient topicDisco;

    // Misc
    private SignedTopicAdvertisement kmcCommunicationTopic;

    private EventProducer producer;
    private EventConsumer consumer;
    private ClientService clientService;

    private class Synclock {

        public synchronized void lock(long timeout) {
            try {
                wait(timeout);
            } catch (InterruptedException e) {}
        }

        public synchronized void unlock() {
            notifyAll();
        }
    }

    Synclock securityTokenRequestLock = new Synclock();

    private boolean waiting;

    private SecureTopicKeyResponse resp;

    private String brokerHost = "";
    private String brokerPort = "";
    private String brokerConnProtocol = "";

    public KMCClient(Certificate cer, PrivateKey priv, PublicKey rootPub,
            String kmcTopic, String configFilePath, String host, String port,
            String prot) {
        init(cer, priv, rootPub, kmcTopic, configFilePath, host, port, prot);
    }

    public KMCClient(Certificate cer, PrivateKey priv, PublicKey rootPub,
            String kmcTopic, String configFilePath) {
        init(cer, priv, rootPub, kmcTopic, configFilePath, null, null, null);
    }

    private void init(Certificate cer, PrivateKey priv, PublicKey rootPub,
            String kmcTopic, String configFilePath, String host, String port,
            String prot) {
        myCertificate = cer;
        myPrivateKey = priv;
        //        rootCAPublicKey = rootPub;
        kmcClientTopic = kmcTopic;

        brokerHost = host;
        brokerPort = port;
        brokerConnProtocol = prot;

        init(configFilePath);
    }

    /**
     * Publishes / Subscribes to the specified topic UUID. this UUID represents
     * the topic over which interested clients can contact this KMC
     * 
     * @param topicUUID -
     *            String representing the string topic
     * @return TRUE if successful, FALSE if an error occurs
     */
    private boolean initTopicCommunication(String topicUUID) {
        try {
            // Subscribe
            Profile profile = clientService.createProfile(
                    TemplateProfileAndSynopsisTypes.STRING, topicUUID);
            consumer.subscribeTo(profile);
        } catch (ServiceException e) {
            log.error("", e);
            return false;
        }
        return true;
    }

    /**
     * Initializes the KMCClient. This process involves finding a broker to
     * connect to (uses Broker Discovery mechanisms), Using topic discovery to
     * find a suitable KMC, and subscribing to appropriate topics.
     * 
     * @return TRUE if successful, FALSE otherwise
     */
    public boolean init(String configFilePath) {

        // Step 1: Uses Broker Discovery to locate a suitable broker and connect
        // to it.

        // IMP: Might conflict w/ other ppl's entity IDs Maintain constant
        // for now
        int entityId = (new Random()).nextInt();

        if (brokerHost == null) {
            // USE Broker Discovery to locate broker !
            // BROKER DISCOVERY:---

            BrokerDiscoveryHelper bdh = new BrokerDiscoveryHelper(
                    (new Random()).nextInt(), configFilePath, 5000, 1, 1);

            BrokerDiscoveryResponse[] ba = bdh.discover(
                    new String[] { "niotcp" }, "**", "*");

            Properties props = (Properties) ba[0].brokerInfo
                    .getSupportedTransportProtocols().get("niotcp");

            brokerHost = ba[0].brokerInfo.getHostname();
            brokerPort = props.getProperty("NIOTCPServerPort");
            brokerConnProtocol = "niotcp";
        }
        // ---

        Properties prop = new Properties();
        prop.put("hostname", brokerHost);
        prop.put("portnum", brokerPort);

        log.info("Connecting using " + brokerConnProtocol + " @ "
                + prop.get("hostname") + ":" + prop.get("portnum")
                + " EntityID[" + entityId + "]");

        try {
            clientService = SessionService.getClientService(entityId,
                    configFilePath);
            clientService.initializeBrokerCommunications(prop,
                    brokerConnProtocol);

            // Create event Subscriber
            consumer = clientService.createEventConsumer(this);
            // Create a producer
            producer = clientService.createEventProducer();
            producer.setSuppressRedistributionToSource(true);

        } catch (ServiceException e) {
            log.error("Error connecting to broker ! ", e);
            return false;
        }

        topicDisco = new TopicDiscoveryClient(entityId, configFilePath,
                myCertificate, myPrivateKey, brokerHost, brokerPort,
                brokerConnProtocol);

        // Step 2: Uses Topic disocvery to locate the topic over which the KMC
        // communicates.
        // Once this is discovered we select a KMC from the list and request its
        // certificate. This certificate will be used for fututre communications
        // with the KMC.
        // TODO: Select 1 among multiple KMCs..
        // Currently only 1 matching topic returned...
        SignedTopicAdvertisement[] resps = topicDisco.discover(
                Topics.MATCHING_REGEX, "^/services/KMC*", null, 1);
        if (resps.length > 0) {
            kmcCommunicationTopic = resps[0];
            if (!kmcCommunicationTopic.verify()) {
                log
                        .error("Could not verify Topic Advertisement signature ! Aborting...");
                return false;
            } else {
                log.info("KMCComm uuid: "
                        + kmcCommunicationTopic.getTopicAd().getUuid());
            }
        } else {
            log.error("Could not locate KMC !");
            return false;
        }

        if (!initTopicCommunication(kmcClientTopic)) {
            log.error("Could not subscribe to kmcClientTopic: "
                    + kmcClientTopic);
            return false;
        }

        topicDisco.close();

        return true;
    }

    public void close() {
        try {
            clientService.closeBrokerConnection();
            clientService.terminateServices();
        } catch (ServiceException e) {
            log.error("", e);
        }
    }

    /**
     * Registers a topic with the kmc to be a secure topic.
     * 
     * @param pubs -
     *            A vector of allowed publishers to this topic
     * @param subs -
     *            A vector of allowed subscribers to this topic
     * @param sta -
     *            The signed topic advertisement. The signature will be verified
     *            with the public key.
     * @param creds -
     *            The Certificate of the entity (owner of the topic)
     * @param validUntil -
     *            Validity period of the secure topic
     * @param keyGenAlgo -
     *            The Secret key generation algorithm to use (default AES)
     * @param keySize -
     *            Size of the key (default 1024)
     * @param timeout -
     *            Time for which to wait for the final response from the KMC
     * @return Secrect topic key and tokenm if successful, NULL otherwise
     */
    public SecureTopicKeyResponse registerTopic(Hashtable pubs, Hashtable subs,
            SignedTopicAdvertisement sta, Certificate creds,
            Calendar validUntil, String keyGenAlgo, int keySize, int timeout) {

        if (keyGenAlgo == null) {
            keyGenAlgo = "AES";
        }

        if (keySize == -1) {
            keySize = 192;
        }

        SecureTopicKeyCreateUpdateRequest req = new SecureTopicKeyCreateUpdateRequest(
                pubs, subs, sta, creds, validUntil, keyGenAlgo, keySize,
                kmcClientTopic);

        SignedRequest sReq = new SignedRequest(req.getBytes(), myCertificate);
        sReq.generateSignedDigest(myPrivateKey);

        EncryptedMessagePayload emp = EncryptedMessagePayload.generate(sReq
                .getBytes(), kmcCommunicationTopic.getOwnerCertificate()
                .getPublicKey());

        KMCMessage msg = new KMCMessage(
                KMCMessage.KmcMSG_SECURE_TOPIC_KEY_CREATE_MODIFY, emp
                        .getBytes());
        send(kmcCommunicationTopic.getTopicAd().getUuid(), msg.getBytes());

        // Now this thread of execution waits until we get a verification
        // request from KMC. We decode this random UUID with out private key and
        // then encode it witht he kmc's public key and resend it back. The kmc
        // will then follow-up with a Secure topic key and security token.
        // We also employ a time out to discard responses which don't meet QoS
        // restrictions

        resp = null;
        waiting = true;
        securityTokenRequestLock.lock(timeout);
        waiting = false;
        return resp;
    }

    /**
     * Request a topic key to publish or subscribe to a topic owned by some
     * other entity
     * 
     * @param synopsis -
     *            The topic synopsis in question
     * @param creds
     * @param reqRights
     * @param timeout
     * @return
     */
    public SecureTopicKeyResponse requestTopicKey(String synopsis,
            Certificate creds, TopicRights reqRights, int timeout) {

        SecureTopicKeyRequest req = new SecureTopicKeyRequest(synopsis, creds,
                reqRights, kmcClientTopic);

        EncryptedMessagePayload emp = EncryptedMessagePayload.generate(req
                .getBytes(), kmcCommunicationTopic.getOwnerCertificate()
                .getPublicKey());

        KMCMessage msg = new KMCMessage(
                KMCMessage.KmcMSG_SECURE_TOPIC_KEY_REQUEST, emp.getBytes());

        send(kmcCommunicationTopic.getTopicAd().getUuid(), msg.getBytes());

        // Now this thread of execution waits until we get a response from the
        // selected KMC.
        // We also employ a time out to discard responses which don't meet QoS
        // restrictions

        resp = null;
        waiting = true;
        securityTokenRequestLock.lock(timeout);
        waiting = false;
        return resp;
    }

    /**
     * Send the specified message on specified topic
     * 
     * @param topicName -
     *            The topic on which to send the message
     * @param request -
     *            The message to be sent
     */
    public void send(String topicName, byte[] request) {
        producer.generateEventIdentifier(true);
        producer.setTemplateId(12345);
        producer.setDisableTimestamp(false);
        try {
            NBEvent nbEvent = producer.generateEvent(
                    TemplateProfileAndSynopsisTypes.STRING, topicName, request);
            producer.publishEvent(nbEvent);
            log
                    .info("sent " + request.length + " bytes ON <" + topicName
                            + ">");
        } catch (ServiceException e) {
            log.error("", e);
        }
    }

    public void onEvent(NBEvent nbEvent) {
        byte[] payload = nbEvent.getContentPayload();

        byte[] objBytes = new byte[payload.length - 1];
        System.arraycopy(payload, 1, objBytes, 0, objBytes.length);

        if (payload[0] == MSG_SECURITY_TOKEN) {
            // PROCESS reciept of security token

            if (waiting) {
                // Decode the uuid
                EncryptedMessagePayload encResp = EncryptedMessagePayload
                        .createObjectFromBytes(objBytes);
                byte[] respBytes = encResp.getDecryptedPayload(myPrivateKey);

                // Ok, store the response and notify waiting threads !!

                resp = SecureTopicKeyResponse.createObjectFromBytes(respBytes);
                securityTokenRequestLock.unlock();
            } else {
                log.debug("Stray Security token recieved ! "
                        + "Probably a TIMEOUT occurred previously !! "
                        + "Discarding...");
            }

        } else if (payload[0] == MSG_ERROR) {
            // PROCESS Error message...

            if (waiting) {
                // Decode the message
                String s = new String(objBytes);

                // Ok, store the response and notify waiting threads !!
                log.error(s);
            } else {
                log.debug("Stray Security token recieved ! "
                        + "Probably a TIMEOUT occurred previously !! "
                        + "Discarding...");
            }
        } else {
            // Garbled header !
            log.error("Garbled header ! Discarding...");
        }
    }
}
